// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/expression/complex_view_types.h>
#include <pollux/functions/udf.h>
#include <pollux/type/floating_point_util.h>

namespace kumo::pollux::functions {
    /// Fast path for constant primitive type keys: map_subset(m, array[1, 2, 3]).
    template<typename TExec, typename Key>
    struct MapSubsetPrimitiveFunction {
        POLLUX_DEFINE_FUNCTION_TYPES(TExec);

        void initialize(
            const std::vector<TypePtr> & /*inputTypes*/,
            const core::QueryConfig & /*config*/,
            const arg_type<Map<Key, Generic<T1> > > * /*inputMap*/,
            const arg_type<Array<Key> > *keys) {
            if (keys != nullptr) {
                constantSearchKeys_ = true;
                initializeSearchKeys(*keys);
            }
        }

        void call(
            out_type<Map<Key, Generic<T1> > > &out,
            const arg_type<Map<Key, Generic<T1> > > &inputMap,
            const arg_type<Array<Key> > &keys) {
            if (keys.empty()) {
                return;
            }

            if (!constantSearchKeys_) {
                searchKeys_.clear();
                initializeSearchKeys(keys);
            }

            if (searchKeys_.empty()) {
                return;
            }

            auto toFind = searchKeys_.size();

            for (const auto &entry: inputMap) {
                if (!searchKeys_.contains(entry.first)) {
                    continue;
                }

                if (!entry.second.has_value()) {
                    auto &keyWriter = out.add_null();
                    keyWriter = entry.first;
                } else {
                    auto [keyWriter, valueWriter] = out.add_item();
                    keyWriter = entry.first;
                    valueWriter.copy_from(entry.second.value());
                }

                --toFind;
                if (toFind == 0) {
                    break;
                }
            }
        }

    private:
        void initializeSearchKeys(const arg_type<Array<Key> > &keys) {
            for (const auto &key: keys.skipNulls()) {
                searchKeys_.emplace(key);
            }
        }

        bool constantSearchKeys_{false};
        util::floating_point::HashSetNaNAware<arg_type<Key> > searchKeys_;
    };

    /// Fast path for constant string keys: map_subset(m, array['a', 'b', 'c']).
    template<typename TExec>
    struct MapSubsetVarcharFunction {
        POLLUX_DEFINE_FUNCTION_TYPES(TExec);

        void initialize(
            const std::vector<TypePtr> & /*inputTypes*/,
            const core::QueryConfig & /*config*/,
            const arg_type<Map<Varchar, Generic<T1> > > * /*inputMap*/,
            const arg_type<Array<Varchar> > *keys) {
            if (keys != nullptr) {
                constantSearchKeys_ = true;

                searchKeyStrings_.reserve(keys->size());
                for (const auto &key: keys->skipNulls()) {
                    if (key.isInline()) {
                        searchKeys_.emplace(key);
                    } else if (!searchKeys_.contains(key)) {
                        searchKeyStrings_.push_back(key.str());
                        searchKeys_.emplace(StringView(searchKeyStrings_.back()));
                    }
                }
            }
        }

        void call(
            out_type<Map<Varchar, Generic<T1> > > &out,
            const arg_type<Map<Varchar, Generic<T1> > > &inputMap,
            const arg_type<Array<Varchar> > &keys) {
            if (keys.empty()) {
                return;
            }

            if (!constantSearchKeys_) {
                searchKeys_.clear();
                for (const auto &key: keys.skipNulls()) {
                    searchKeys_.emplace(key);
                }
            }

            if (searchKeys_.empty()) {
                return;
            }

            auto toFind = searchKeys_.size();

            for (const auto &entry: inputMap) {
                if (!searchKeys_.contains(entry.first)) {
                    continue;
                }

                if (!entry.second.has_value()) {
                    auto &keyWriter = out.add_null();
                    keyWriter.copy_from(entry.first);
                } else {
                    auto [keyWriter, valueWriter] = out.add_item();
                    keyWriter.copy_from(entry.first);
                    valueWriter.copy_from(entry.second.value());
                }

                --toFind;
                if (toFind == 0) {
                    break;
                }
            }
        }

    private:
        bool constantSearchKeys_{false};
        melon::F14FastSet<StringView> searchKeys_;
        std::vector<std::string> searchKeyStrings_;
    };

    /// Generic implementation. Doesn't provide an optimization for constant search
    /// keys.
    template<typename TExec>
    struct MapSubsetFunction {
        POLLUX_DEFINE_FUNCTION_TYPES(TExec);

        void call(
            out_type<Map<Generic<T1>, Generic<T2> > > &out,
            const arg_type<Map<Generic<T1>, Generic<T2> > > &inputMap,
            const arg_type<Array<Generic<T1> > > &keys) {
            if (keys.empty()) {
                return;
            }

            // TODO Figure out how to implement fast path for constant search keys.
            // Just implementing 'initialize' and populating 'searchKeys_' there doesn't
            // work because GenericView's go out of scope at the end of the method.

            searchKeys_.clear();
            for (const auto &key: keys.skipNulls()) {
                searchKeys_.emplace(key);
            }

            if (searchKeys_.empty()) {
                return;
            }

            auto toFind = searchKeys_.size();

            for (const auto &entry: inputMap) {
                if (!searchKeys_.contains(entry.first)) {
                    continue;
                }

                if (!entry.second.has_value()) {
                    auto &keyWriter = out.add_null();
                    keyWriter.copy_from(entry.first);
                } else {
                    auto [keyWriter, valueWriter] = out.add_item();
                    keyWriter.copy_from(entry.first);
                    valueWriter.copy_from(entry.second.value());
                }

                --toFind;
                if (toFind == 0) {
                    break;
                }
            }
        }

    private:
        melon::F14FastSet<exec::GenericView> searchKeys_;
    };
} // namespace kumo::pollux::functions
